from googleapiclient.discovery import build
import os
import time
import requests
import socks
import socket
import argparse
from urllib3.contrib.socks import SOCKSProxyManager
from dotenv import load_dotenv
import csv
from datetime import datetime, timedelta

class YouTubeAnalyzer:
    def __init__(self, api_key, use_proxy=False, proxy_url=None):
        self.start_time = time.time()  # 记录开始时间
        self.api_key = api_key
        self.max_retries = 5
        self.retry_delay = 10
        
        # 设置代理
        if use_proxy and proxy_url:
            proxy_parts = proxy_url.replace('socks5://', '').split(':')
            socks.set_default_proxy(socks.SOCKS5, proxy_parts[0], int(proxy_parts[1]))
            socket.socket = socks.socksocket
        
        self.youtube = build('youtube', 'v3', developerKey=self.api_key)

    def format_time_duration(self, seconds):
        """格式化执行时间"""
        duration = timedelta(seconds=int(seconds))
        hours = duration.seconds // 3600
        minutes = (duration.seconds % 3600) // 60
        seconds = duration.seconds % 60
        
        time_parts = []
        if hours > 0:
            time_parts.append(f"{hours}小时")
        if minutes > 0:
            time_parts.append(f"{minutes}分钟")
        if seconds > 0 or not time_parts:
            time_parts.append(f"{seconds}秒")
            
        return "".join(time_parts)

    def _make_request_with_retry(self, request):
        """带重试机制的请求"""
        for attempt in range(self.max_retries):
            try:
                return request.execute()
            except Exception as e:
                if attempt == self.max_retries - 1:
                    print(f"最终请求失败: {str(e)}")
                    return None
                print(f"请求失败，{self.retry_delay}秒后重试: {str(e)}")
                time.sleep(self.retry_delay)

    def get_video_info(self, video_url):
        """从视频URL获取视频信息"""
        try:
            # 从URL中提取视频ID
            video_id = video_url.split('/shorts/')[1].split('?')[0]
            
            request = self.youtube.videos().list(
                part='snippet',
                id=video_id
            )
            response = self._make_request_with_retry(request)
            
            if response and response.get('items'):
                return response['items'][0]['snippet']['channelId']
            return None
        except Exception as e:
            print(f"获取视频信息时出错: {str(e)}")
            return None

    def get_channel_videos_count(self, channel_id):
        """获取频道的短视频数量"""
        try:
            # 获取短视频数量
            shorts_request = self.youtube.search().list(
                part='id',
                channelId=channel_id,
                type='video',
                videoDuration='short',
                maxResults=50  # API限制，每次最多50条
            )
            shorts_response = self._make_request_with_retry(shorts_request)
            shorts_total = shorts_response.get('pageInfo', {}).get('totalResults', 0)

            return shorts_total
        except Exception as e:
            print(f"获取视频数量时出错: {str(e)}")
            return 0

    def get_channel_info(self, channel_id):
        """获取频道信息"""
        try:
            request = self.youtube.channels().list(
                part='snippet,statistics',
                id=channel_id
            )
            response = self._make_request_with_retry(request)
            
            if response and response.get('items'):
                channel = response['items'][0]
                
                # 只获取短视频数量
                shorts_count = self.get_channel_videos_count(channel_id)
                
                return {
                    'channel_url': f"https://www.youtube.com/channel/{channel_id}",
                    'created_at': channel['snippet']['publishedAt'],
                    'subscriber_count': channel['statistics'].get('subscriberCount', 'N/A'),
                    'view_count': channel['statistics'].get('viewCount', 'N/A'),
                    'shorts_count': shorts_count
                }
            return None
        except Exception as e:
            print(f"获取频道信息时出错: {str(e)}")
            return None

    def process_video_urls(self, input_file):
        """处理输入文件中的视频URL"""
        phase_start_time = time.time()  # 记录阶段开始时间
        channels_info = []
        processed_channels = set()  # 用于去重
        
        # 尝试不同的编码方式
        encodings = ['utf-8-sig', 'utf-8', 'gbk', 'gb2312', 'iso-8859-1']
        
        for encoding in encodings:
            try:
                with open(input_file, 'r', encoding=encoding) as f:
                    reader = csv.reader(f)
                    next(reader)  # 跳过表头
                    
                    for row in reader:
                        if not row:
                            continue
                            
                        video_url = row[0]
                        print(f"\n处理视频: {video_url}")
                        
                        channel_id = self.get_video_info(video_url)
                        if not channel_id or channel_id in processed_channels:
                            continue
                        
                        processed_channels.add(channel_id)
                        print(f"获取频道 {channel_id} 的信息...")
                        
                        channel_info = self.get_channel_info(channel_id)
                        if channel_info:
                            channels_info.append(channel_info)
                        
                        time.sleep(1)  # 添加延迟避免请求过快
                    
                    # 如果成功读取，跳出循环
                    break
                        
            except UnicodeDecodeError:
                # 如果是最后一种编码方式仍然失败，则报错
                if encoding == encodings[-1]:
                    print(f"无法读取文件，请检查文件编码")
                    return []
                continue
            except Exception as e:
                print(f"处理输入文件时出错: {str(e)}")
                return []
                
        # 计算并显示处理阶段耗时
        phase_duration = time.time() - phase_start_time
        print(f"\n处理视频链接耗时: {self.format_time_duration(phase_duration)}")
        
        return channels_info

    def save_channel_info_csv(self, tag, channels_info):
        """保存频道信息到CSV文件"""
        phase_start_time = time.time()  # 记录阶段开始时间
        
        timestamp = time.strftime("%Y%m%d%H%M%S")
        filename = f"{tag}-{timestamp}.csv"
        
        headers = [
            '频道链接', 
            '注册时间', 
            '订阅数量', 
            '观看次数', 
            '短视频数量'
        ]
        
        try:
            with open(filename, 'w', newline='', encoding='utf-8-sig') as f:
                writer = csv.writer(f)
                writer.writerow(headers)
                
                for info in channels_info:
                    if info:
                        writer.writerow([
                            info['channel_url'],
                            info['created_at'],
                            info['subscriber_count'],
                            info['view_count'],
                            info['shorts_count']
                        ])
            
            # 计算并显示保存阶段耗时
            phase_duration = time.time() - phase_start_time
            print(f"保存文件耗时: {self.format_time_duration(phase_duration)}")
            print(f"\n结果已保存到CSV文件: {filename}")
            
        except Exception as e:
            print(f"保存CSV文件时出错: {str(e)}")

    def __del__(self):
        """清理资源并显示总耗时"""
        # 计算并显示总耗时
        if hasattr(self, 'start_time'):
            total_duration = time.time() - self.start_time
            print(f"\n程序总耗时: {self.format_time_duration(total_duration)}")

def main():
    print(f"程序开始执行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='处理YouTube短视频链接并获取频道信息')
    parser.add_argument('input_file', help='输入的CSV文件路径')
    parser.add_argument('tag', help='标签名称（用于输出文件名）')
    args = parser.parse_args()
    
    # 加载环境变量
    load_dotenv()
    
    # 从环境变量获取配置
    API_KEY = os.getenv('YOUTUBE_API_KEY')
    USE_PROXY = os.getenv('USE_PROXY', 'False').lower() == 'true'
    PROXY_URL = os.getenv('PROXY_URL')
    
    if not API_KEY:
        raise ValueError("请在 .env 文件中设置 YOUTUBE_API_KEY")
    
    analyzer = YouTubeAnalyzer(
        api_key=API_KEY,
        use_proxy=USE_PROXY,
        proxy_url=PROXY_URL if USE_PROXY else None
    )
    
    print(f"开始处理文件: {args.input_file}")
    
    # 处理视频链接并获取频道信息
    channels_info = analyzer.process_video_urls(args.input_file)
    
    if not channels_info:
        print("未获取到任何频道信息")
        return
    
    # 保存结果
    analyzer.save_channel_info_csv(args.tag, channels_info)

if __name__ == "__main__":
    main() 